#include <iostream>
#include <vector>

#include <gtest/gtest.h>
#include "op_log.h"
#define private public

#include "kernel_run_context_facker.h"

#include "fusion_ops.h"
#include "array_ops.h"
#include "op_tiling/op_tiling_util.h"
#include "common/utils/ut_op_util.h"
#include "common_unittest.h"
#include "exe_graph/runtime/storage_format.h"
#include "exe_graph/runtime/storage_shape.h"
#include "exe_graph/runtime/tiling_parse_context.h"
#include "test_cube_util.h"
#include "runtime2_util.h"
#include "platform/platform_info.h"
#include "register/op_def_registry.h"
#include "tiling/tiling_api.h"

using namespace std;

class PromptFlashAttentionTiling : public testing::Test
{
protected:
  static void SetUpTestCase()
  {
    std::cout << "PromptFlashAttentionTiling SetUp" << std::endl;
  }

  static void TearDownTestCase()
  {
    std::cout << "PromptFlashAttentionTiling TearDown" << std::endl;
  }
};

const map<string, string> kPlatformInfo{{"Ascend310P", R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
                          "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "cube_core_cnt": 20, "vector_core_cnt": 40, "core_type_list": "CubeCore,VectorCore"}
                          })"},
                                        {"Ascend910B", R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
                          "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "cube_core_cnt": 24, "vector_core_cnt": 48, "core_type_list": "CubeCore,VectorCore"}
                          })"},
                                        {"Ascend310P3", R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
                          "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 262144, "L2_SIZE": 16777216, "L1_SIZE": 1048576,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 262144,
                          "cube_core_cnt": 8, "vector_core_cnt": 7, "core_type_list": "CubeCore,VectorCore"}
                          })"}};

static string TilingData2Str(const gert::TilingData *tiling_data)
{
  auto data = tiling_data->GetData();
  string result;
  for (size_t i = 0; i < tiling_data->GetDataSize(); i += sizeof(int32_t))
  {
    result += std::to_string((reinterpret_cast<const int32_t *>(tiling_data->GetData())[i / sizeof(int32_t)]));
    result += " ";
  }

  return result;
}

using namespace ge;
using namespace ut_util;

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_1)
{

  gert::StorageShape query_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape key_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape value_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 5120}, {1, 1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_2)
{

  gert::StorageShape query_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape key_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape value_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_bf16)
{

  gert::StorageShape query_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape key_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape value_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 0);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_3)
{

  gert::StorageShape query_shape = {{2, 1024, 5120}, {2, 1024, 5120}};
  gert::StorageShape key_shape = {{2, 1024, 5120}, {2, 1024, 5120}};
  gert::StorageShape value_shape = {{2, 1024, 5120}, {2, 1024, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  std::string input_layout = "BSH";
  int64_t num_key_value_heads = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_int8_in_fp16_out) {

  gert::StorageShape query_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape key_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  // gert::StorageShape pse_shift = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape attenmask = {{1, 40, 1000, 1000}, {1, 40, 1000, 1000}};
  gert::StorageShape actual_seq_lengths = {{1}, {1}};
  gert::StorageShape actual_seq_lengths_kv = {{1}, {1}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  std::string input_layout = "BNSD";
  int64_t num_key_value_heads = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(10, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &attenmask, &actual_seq_lengths, &actual_seq_lengths_kv, &deq_scale1, &quant_scale1, &deq_scale2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000001217);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_int8_in_int8_out) {
  gert::StorageShape query_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape key_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  // gert::StorageShape pse_shift = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape attenmask = {{1, 40, 1000, 1000}, {1, 40, 1000, 1000}};
  gert::StorageShape actual_seq_lengths = {{1}, {1}};
  gert::StorageShape actual_seq_lengths_kv = {{1}, {1}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  gert::StorageShape quant_scale2 = {{1}, {1}};
  gert::StorageShape quant_offset2 = {{1}, {1}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  std::string input_layout = "BNSD";
  int64_t num_key_value_heads = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &attenmask, &actual_seq_lengths, &actual_seq_lengths_kv, &deq_scale1, &quant_scale1, &deq_scale2, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000021217);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_4) {
  gert::StorageShape query_shape = {{2, 1000, 5120}, {2, 1000, 5120}};
  gert::StorageShape key_shape = {{2, 1000, 5120}, {2, 1000, 5120}};
  gert::StorageShape value_shape = {{2, 1000, 5120}, {2, 1000, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  std::string input_layout = "BSH";
  int64_t num_key_value_heads = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 5120}, {1, 1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}
TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_5)
{

  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000001012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_6)
{

  gert::StorageShape query_shape = {{40, 1024, 128}, {40, 1024, 128}};
  gert::StorageShape key_shape = {{40, 1024, 128}, {40, 1024, 128}};
  gert::StorageShape value_shape = {{40, 1024, 128}, {40, 1024, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "NSD";

  std::vector<gert::StorageShape> output_shapes(1, {{40, 1024, 128}, {40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000001012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}
TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_7)
{

  gert::StorageShape query_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape key_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape value_shape = {{256, 5120}, {256, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{256, 5120}, {256, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_SH_attenMask)
{

  gert::StorageShape query_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape key_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape value_shape = {{256, 5120}, {256, 5120}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 256, 128}, {1, 40, 256, 128}};
  gert::StorageShape atten_mask_shape = {{256, 256}, {256, 256}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{256, 5120}, {256, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_SH_attenMask_sparse_mode1)
{

  gert::StorageShape query_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape key_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape value_shape = {{256, 5120}, {256, 5120}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 256, 128}, {1, 40, 256, 128}};
  gert::StorageShape atten_mask_shape = {{256, 256}, {256, 256}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  int64_t sparse_mode = 1;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{256, 5120}, {256, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_SH_attenMask_sparse_mode3)
{

  gert::StorageShape query_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape key_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape value_shape = {{256, 5120}, {256, 5120}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 256, 128}, {1, 40, 256, 128}};
  gert::StorageShape atten_mask_shape = {{2048, 2048}, {2048, 2048}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  int64_t sparse_mode = 3;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{256, 5120}, {256, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_SH_attenMask_sparse_mode4)
{

  gert::StorageShape query_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape key_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape value_shape = {{256, 5120}, {256, 5120}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 256, 128}, {1, 40, 256, 128}};
  gert::StorageShape atten_mask_shape = {{2048, 2048}, {2048, 2048}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  int64_t sparse_mode = 4;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{256, 5120}, {256, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_BNSD_attenMask_sparse_mode4_case0)
{

  gert::StorageShape query_shape = {{1,8,2048, 128}, {1,8,2048, 128}};
  gert::StorageShape key_shape = {{1,8,2048, 128}, {1,8,2048, 128}};
  gert::StorageShape value_shape = {{1,8,2048, 128}, {1,8,2048, 128}};
  gert::StorageShape atten_mask_shape = {{2048, 2048}, {2048, 2048}};
  int64_t num_heads = 8;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 8;
  int64_t sparse_mode = 4;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1,8,2048, 128}, {1,8,2048, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 50}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_BNSD_attenMask_sparse_mode4_case1)
{
  gert::StorageShape query_shape = {{1,12,16384,128}, {1,12,16384,128}};
  gert::StorageShape key_shape = {{1,12,16384,128}, {1,12,16384,128}};
  gert::StorageShape value_shape = {{1,12,16384,128}, {1,12,16384,128}};
  gert::StorageShape atten_mask_shape = {{2048, 2048}, {2048, 2048}};
  int64_t num_heads = 12;
  float scale_value = 1.0f;
  int64_t pre_tokens = 16 * 1024;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 12;
  int64_t sparse_mode = 4;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1,12,16384,128}, {1,12,16384,128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 50}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_SH_attenMaskUint8)
{

  gert::StorageShape query_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape key_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape value_shape = {{256, 5120}, {256, 5120}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 256, 128}, {1, 40, 256, 128}};
  gert::StorageShape atten_mask_shape = {{256, 256}, {256, 256}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{256, 5120}, {256, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_UINT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_BSH_attenMask)
{

  gert::StorageShape query_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape key_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape value_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape atten_mask_shape = {{1, 1000, 1000}, {1, 1000, 1000}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 5120}, {1, 1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_8)
{

  gert::StorageShape query_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape key_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000001012);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_high_precision0)
{
  gert::StorageShape query_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape key_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape value_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";
  int64_t sparse_mode = 0;
  int64_t inner_precise = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)},
                                {"inner_precise", ge::AnyValue::CreateFrom<int64_t>(inner_precise)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101612);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_sparsemode_0_negative_nextoken)
{
  gert::StorageShape query_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape key_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape value_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = -1000;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";
  int64_t sparse_mode = 0;
  int64_t inner_precise = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)},
                                {"inner_precise", ge::AnyValue::CreateFrom<int64_t>(inner_precise)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions);
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 1000000000000101612);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_high_precision1)
{
  gert::StorageShape query_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape key_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape value_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  int64_t num_heads = 20;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";
  int64_t sparse_mode = 0;
  int64_t inner_precise = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)},
                                {"inner_precise", ge::AnyValue::CreateFrom<int64_t>(inner_precise)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101612);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_high_precision2)
{
  gert::StorageShape query_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape key_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape value_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  int64_t num_heads = 80;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";
  int64_t sparse_mode = 0;
  int64_t inner_precise = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)},
                                {"inner_precise", ge::AnyValue::CreateFrom<int64_t>(inner_precise)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101612);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_high_precision3)
{
  gert::StorageShape query_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape key_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  gert::StorageShape value_shape = {{1, 1024, 5120}, {1, 1024, 5120}};
  int64_t num_heads = 160;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";
  int64_t sparse_mode = 0;
  int64_t inner_precise = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1024, 5120}, {1, 1024, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)},
                                {"inner_precise", ge::AnyValue::CreateFrom<int64_t>(inner_precise)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000101612);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_d_unalign_1) {

  gert::StorageShape query_shape = {{1, 40, 1000, 127}, {1, 40, 1000, 127}};
  gert::StorageShape key_shape = {{1, 40, 1000, 127}, {1, 40, 1000, 127}};
  gert::StorageShape value_shape = {{1, 40, 1000, 127}, {1, 40, 1000, 127}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 127}, {1, 40, 1000, 127}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000000016);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_empty_q)
{
  gert::StorageShape query_shape = {{1, 40, 0, 128}, {1, 40, 0, 128}};
  gert::StorageShape key_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // return failed for empty q
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_empty_k_bsnd)
{
  gert::StorageShape query_shape = {{1, 1000, 40, 128}, {1, 1000, 40, 128}};
  gert::StorageShape key_shape = {{1, 0, 40, 128}, {1, 0, 40, 128}};
  gert::StorageShape value_shape = {{1, 1000, 40, 128}, {1, 1000, 40, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSND";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 40, 128}, {1, 1000, 40, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000000020);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_empty_k_bnsd)
{
  gert::StorageShape query_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape key_shape = {{1, 40, 0, 128}, {1, 40, 0, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000000020);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_empty_k_bsh)
{
  gert::StorageShape query_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape key_shape = {{1, 0, 5120}, {1, 0, 5120}};
  gert::StorageShape value_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 5120}, {1, 1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000000020);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_empty_k_sh)
{
  gert::StorageShape query_shape = {{1000, 5120}, {1000, 5120}};
  gert::StorageShape key_shape = {{0, 5120}, {0, 5120}};
  gert::StorageShape value_shape = {{1000, 5120}, {1000, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{1000, 5120}, {1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000000020);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_empty_atten_mask)
{
  gert::StorageShape query_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape key_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 0, 128}, {1, 40, 0, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 16);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, gqa_test0)
{

  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 8, 1024, 128}, {1, 8, 1024, 128}};
  gert::StorageShape value_shape = {{1, 8, 1024, 128}, {1, 8, 1024, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 8;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 15);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_sparse_mode_and_empty_atten_mask)
{
  gert::StorageShape query_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape key_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 0, 128}, {1, 40, 0, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  int64_t sparse_mode = 0; // SPARSE_MODE_NO_MASK
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 16);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_sparse_mode_and_empty_atten_mask_2)
{
  gert::StorageShape query_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape key_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape value_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape pse_shift_shape = {{1, 40, 1000, 128}, {1, 40, 1000, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 0, 128}, {1, 40, 0, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  int64_t sparse_mode = 2; // SPARSE_MODE_LEFT_UP
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1000, 128}, {1, 40, 1000, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, &pse_shift_shape, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)},
                                {"sparse_mode", ge::AnyValue::CreateFrom<int64_t>(sparse_mode)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 16);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, gqa_test1_invalid_kvheads)
{

  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 8, 1024, 128}, {1, 8, 1024, 128}};
  gert::StorageShape value_shape = {{1, 8, 1024, 128}, {1, 8, 1024, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 7;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // todo check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 15);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}
/*
TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_9) {

  gert::StorageShape query_shape = {{1, 2048, 640}, {1, 2048, 640}};
  gert::StorageShape key_shape = {{1, 2048, 640}, {1, 2048, 640}};
  gert::StorageShape value_shape = {{1, 2048, 640}, {1, 2048, 640}};
  int64_t num_heads = 5;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  std::string input_layout = "BSH";
  int64_t num_key_value_heads = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 2048, 640}, {1, 2048, 640}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {};
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 12);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

/*
TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_11_kv_diff) {

  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 512, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 512, 128}, {1, 40, 1024, 128}};
  gert::StorageShape pse_shift_shape = {{1, 1, 2048}, {1, 1, 2048}};
  gert::StorageShape atten_mask_shape = {{1, 1, 2048}, {1, 1, 2048}};
  gert::StorageShape actual_seq_lens = {{2}, {2}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 1;
  int64_t next_tokens = 0;
  std::string input_layout = "BNSD";
  int64_t num_key_value_heads = 0;

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {};
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(6, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, &pse_shift_shape, &atten_mask_shape, &actual_seq_lens})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 16);
  auto block_dim = tiling_context->GetBlockDim();
  ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}
*/

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_atten_mask_datatype_is_bool)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000001012);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_fp16_in_int8_out_BNSD)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  gert::StorageShape quant_scale2 = {{1}, {1}};
  gert::StorageShape quant_offset2 = {{1}, {1}};

  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape, nullptr, nullptr, &deq_scale1, &quant_scale1, &deq_scale2, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
//   ASSERT_EQ(tiling_key, 1000000000000021012);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_fp16_in_int8_out_per_channel_case0)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 40}, {1, 40, 1024, 40}};
  gert::StorageShape key_shape = {{1, 40, 1024, 40}, {1, 40, 1024, 40}};
  gert::StorageShape value_shape = {{1, 40, 1024, 40}, {1, 40, 1024, 40}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  gert::StorageShape quant_scale2 = {{40, 40}, {40, 40}};
  gert::StorageShape quant_offset2 = {{40, 40}, {40, 40}};

  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 40}, {1, 40, 1024, 40}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape, nullptr, nullptr, nullptr, nullptr, nullptr, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_fp16_in_int8_out_BSH)
{

  gert::StorageShape query_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape key_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape value_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  gert::StorageShape quant_scale2 = {{1}, {1}};
  gert::StorageShape quant_offset2 = {{1}, {1}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 5120}, {1, 1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, nullptr, nullptr, nullptr, &deq_scale1, &quant_scale1, &deq_scale2, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
//   ASSERT_EQ(tiling_key, 1000000000000121012);
  auto block_dim = tiling_context->GetBlockDim();
//   ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_fp16_in_int8_out_SH_FAILED)
{

  gert::StorageShape query_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape key_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape value_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 5120}, {1, 1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_XY_smallS)
{
  gert::StorageShape query_shape = {{1, 1, 4, 128}, {1, 1, 4, 128}};
  gert::StorageShape key_shape = {{1, 1, 2048, 128}, {1, 1, 2048, 128}};
  gert::StorageShape value_shape = {{1, 1, 2048, 128}, {1, 1, 2048, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 1, 4, 128}, {1, 1, 4, 128}};
  gert::StorageShape atten_mask_shape = {{1, 1, 4, 2048}, {1, 1, 4, 2048}};
  int64_t num_heads = 1;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1, 4, 128}, {1, 1, 4, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  // auto tiling_key = tiling_context->GetTilingKey();
  // ASSERT_EQ(tiling_key, 1015);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_atten_mask_datatype_is_int8)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000001012);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

/*
TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_atten_mask_datatype_is_float32) {
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {};
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, &pse_shift_shape, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 15);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}


TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_atten_mask_datatype_is_float16) {
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {};
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, &pse_shift_shape, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 15);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

*/

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_310P_BSH)
{

  gert::StorageShape query_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape key_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  gert::StorageShape value_shape = {{1, 1000, 5120}, {1, 1000, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 1000, 5120}, {1, 1000, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend310P"}};

  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  compile_info.socShortName = platform_ascendc::SocVersion::ASCEND310P;
  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  tiling_context->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  tiling_context->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  tiling_context->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  tiling_context->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  tiling_func(tiling_context);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_310P_NSD)
{

  gert::StorageShape query_shape = {{40, 1024, 128}, {40, 1024, 128}};
  gert::StorageShape key_shape = {{40, 1024, 128}, {40, 1024, 128}};
  gert::StorageShape value_shape = {{40, 1024, 128}, {40, 1024, 128}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "NSD";

  std::vector<gert::StorageShape> output_shapes(1, {{40, 1024, 128}, {40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend310P"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  compile_info.socShortName = platform_ascendc::SocVersion::ASCEND310P;
  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  tiling_func(tiling_context);
}

// TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_310P_BNSD)
// {

//   gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
//   gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
//   gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
//   int64_t num_heads = 40;
//   float scale_value = 1.0f;
//   int64_t pre_tokens = 65535;
//   int64_t next_tokens = 0;
//   int64_t num_key_value_heads = 0;
//   std::string input_layout = "BNSD";

//   std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

//   std::vector<void *> output_shapes_ref(1);
//   for (size_t i = 0; i < output_shapes.size(); ++i) {
//     output_shapes_ref[i] = &output_shapes[i];
//   }

//   string compile_info_string = R"({
//         "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
//                           "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
//                           "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
//                           "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
//                           "CORE_NUM": 40}
//                           })";
//   map<string, string> soc_infos;
//   map<string, string> aicore_spec;
//   map<string, string> intrinsics;
//   map<string, string> socversions = {{"Short_SoC_version", "Ascend310P"}};
//   GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

//   // platform info
//   fe::PlatFormInfos platform_info;
//   platform_info.Init();
//     // compile info
//   struct PromptFlashAttentionCompileInfo {};
//   PromptFlashAttentionCompileInfo compile_info;

//   // std::map<std::string, std::string> platform_res;
//   // platform_res["ai_core_cnt"] = std::to_string(24); // key
//   // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
//   // platform_info.SetCoreNumByCoreType("AiCore"); // value

//   std::string op_type("PromptFlashAttention");
//   ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
//   auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
//   auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

//   // tilingParseFunc simulate
//   auto kernel_holder = gert::KernelRunContextFaker()
//                     .KernelIONum(2, 1)
//                     .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
//                     .Outputs({&compile_info})
//                     .Build();

//   ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
//   kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
//   kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
//   kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
//   kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

//   ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

//   // tilingFunc simulate
//   auto param = gert::TilingData::CreateCap(4096);
//   auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
//   auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
//   ASSERT_NE(param, nullptr);
//   auto holder = gert::TilingContextFaker()
//                     .NodeIoNum(3, 1)
//                     .IrInstanceNum({1, 1, 1})
//                     .InputShapes({&query_shape, &key_shape, &value_shape})
//                     .OutputShapes(output_shapes_ref)
//                     .CompileInfo(&compile_info)
//                     .PlatformInfo(reinterpret_cast<char *>(&platform_info))
//                     .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
//                     .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
//                     .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
//                     .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
//                     .NodeAttrs({
//                       {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
//                       {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
//                       {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
//                       {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
//                       {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
//                       {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
//                       })
//                     .TilingData(param.get())
//                     .Workspace(ws_size)
//                     .Build();

//   gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
//   ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
//   holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
//   holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
//   holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
//   holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
//   tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
//   // workspaces nullptr return failed
//   tiling_func(tiling_context);
//   // todo check tiling result
// }

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_310P_SH)
{

  gert::StorageShape query_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape key_shape = {{256, 5120}, {256, 5120}};
  gert::StorageShape value_shape = {{256, 5120}, {256, 5120}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "SH";

  std::vector<gert::StorageShape> output_shapes(1, {{256, 5120}, {256, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend310P"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  compile_info.socShortName = platform_ascendc::SocVersion::ASCEND310P;
  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  tiling_func(tiling_context);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_310P_BSND)
{
  gert::StorageShape query_shape = {{1, 2048, 32, 128}, {1, 2048, 32, 128}};
  gert::StorageShape key_shape = {{1, 2048, 32, 128}, {1, 2048, 32, 128}};
  gert::StorageShape value_shape = {{1, 2048, 32, 128}, {1, 2048, 32, 128}};
  int64_t num_heads = 32;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSND";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 2048, 32, 128}, {1, 2048, 32, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"CORE_NUM": 8, "L2_SIZE": 16777216, "L1_SIZE": 1048576, "L0A_SIZE": 65536, "L0B_SIZE": 65536,
        "L0C_SIZE": 262144, "UB_SIZE": 262144, "BT_SIZE": 0, "cube_vector_split_bool": false,
        "ddr_read_rate": 17, "ddr_write_rate": 17, "l2_rate": 114, "l2_read_rate": 114, "l2_write_rate": 114, "l1_to_l0_a_rate": 512,
        "l1_to_l0_b_rate": 256, "l1_to_ub_rate": 128, "l0_c_to_ub_rate": 256, "ub_to_l2_rate": 114,
        "ub_to_ddr_rate": 17, "ub_to_l1_rate": 256, "cube_bandwidth": 0, "vector_bandwidth": 0, "vector_core_cnt": 7}})";

  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend310P"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  compile_info.socShortName = platform_ascendc::SocVersion::ASCEND310P;
  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_310P_BNSD)
{
  gert::StorageShape query_shape = {{1, 32, 2048, 128}, {1, 32, 2048, 128}};
  gert::StorageShape key_shape = {{1, 32, 2048, 128}, {1, 32, 2048, 128}};
  gert::StorageShape value_shape = {{1, 32, 2048, 128}, {1, 32, 2048, 128}};
  int64_t num_heads = 32;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 32, 2048, 128}, {1, 32, 2048, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"CORE_NUM": 8, "L2_SIZE": 16777216, "L1_SIZE": 1048576, "L0A_SIZE": 65536, "L0B_SIZE": 65536,
        "L0C_SIZE": 262144, "UB_SIZE": 262144, "BT_SIZE": 0, "cube_vector_split_bool": false,
        "ddr_read_rate": 17, "ddr_write_rate": 17, "l2_rate": 114, "l2_read_rate": 114, "l2_write_rate": 114, "l1_to_l0_a_rate": 512,
        "l1_to_l0_b_rate": 256, "l1_to_ub_rate": 128, "l0_c_to_ub_rate": 256, "ub_to_l2_rate": 114,
        "ub_to_ddr_rate": 17, "ub_to_l1_rate": 256, "cube_bandwidth": 0, "vector_bandwidth": 0, "vector_core_cnt": 7}})";

  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend310P"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  compile_info.socShortName = platform_ascendc::SocVersion::ASCEND310P;
  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions
  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_1_2) {

  gert::StorageShape query_shape = {{1, 100, 5120}, {1, 100, 5120}};
  gert::StorageShape key_shape = {{1, 100, 5120}, {1, 100, 5120}};
  gert::StorageShape value_shape = {{1, 100, 5120}, {1, 100, 5120}};
  int64_t num_heads = 10;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BSH";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 100, 5120}, {1, 100, 5120}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
    // compile info
  struct PromptFlashAttentionCompileInfo {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  // std::map<std::string, std::string> platform_res;
  // platform_res["ai_core_cnt"] = std::to_string(24); // key
  // platform_info.SetPlatformRes("SoCInfo", platform_res); // label
  // platform_info.SetCoreNumByCoreType("AiCore"); // value

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                    .KernelIONum(2, 1)
                    .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                    .Outputs({&compile_info})
                    .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(3, 1)
                    .IrInstanceNum({1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({
                      {"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                      {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                      {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                      {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                      {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                      {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}
                      })
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  // workspaces nullptr return failed
  //EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // todo check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 0);
  //auto block_dim = tiling_context->GetBlockDim();
  //ASSERT_EQ(block_dim, 40);

  // tiling_data not to check, because of exiting random num from tiling placehold
  // auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
  // ASSERT_EQ(tiling_data_result, tiling_data);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_pse_fp16)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(5, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, &pse_shift_shape, &atten_mask_shape})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
  ASSERT_EQ(tiling_key, 1000000000000000007);
  // auto block_dim = tiling_context->GetBlockDim();
  // ASSERT_EQ(block_dim, 40);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_bf16_in_int8_out_BNSD)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  gert::StorageShape quant_scale2 = {{1}, {1}};
  gert::StorageShape quant_offset2 = {{1}, {1}};

  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape, nullptr, nullptr, &deq_scale1, &quant_scale1, &deq_scale2, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
//   ASSERT_EQ(tiling_key, 1000000000000021012);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_layout_BNSD_BSND_1)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  gert::StorageShape quant_scale2 = {{1}, {1}};
  gert::StorageShape quant_offset2 = {{1}, {1}};

  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "BNSD_BSND";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape, nullptr, nullptr, &deq_scale1, &quant_scale1, &deq_scale2, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
//   ASSERT_EQ(tiling_key, 1000000000000021012);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_invalid_input_layout_1)
{
  gert::StorageShape query_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape key_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape value_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  gert::StorageShape atten_mask_shape = {{1, 40, 1024, 1024}, {1, 40, 1024, 1024}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  gert::StorageShape quant_scale2 = {{1}, {1}};
  gert::StorageShape quant_offset2 = {{1}, {1}};

  int64_t num_heads = 40;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 0;
  std::string input_layout = "";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 40, 1024, 128}, {1, 40, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, &atten_mask_shape, nullptr, nullptr, &deq_scale1, &quant_scale1, &deq_scale2, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
//   ASSERT_EQ(tiling_key, 1000000000000021012);
}

TEST_F(PromptFlashAttentionTiling, prompt_flash_attention_tiling_G_illegal)
{
  gert::StorageShape query_shape = {{1, 100, 1024, 128}, {1, 100, 1024, 128}};
  gert::StorageShape key_shape = {{1, 1, 1024, 128}, {1, 1, 1024, 128}};
  gert::StorageShape value_shape = {{1, 1, 1024, 128}, {1, 1, 1024, 128}};
  // gert::StorageShape pse_shift_shape = {{1, 40, 1024, 128}, {1, 40, 1024, 128}};
  //gert::StorageShape atten_mask_shape = {{1, 100, 1024, 1024}, {1, 100, 1024, 1024}};
  gert::StorageShape deq_scale1 = {{1}, {1}};
  gert::StorageShape quant_scale1 = {{1}, {1}};
  gert::StorageShape deq_scale2 = {{1}, {1}};
  gert::StorageShape quant_scale2 = {{1}, {1}};
  gert::StorageShape quant_offset2 = {{1}, {1}};

  int64_t num_heads = 100;
  float scale_value = 1.0f;
  int64_t pre_tokens = 65535;
  int64_t next_tokens = 0;
  int64_t num_key_value_heads = 1;
  std::string input_layout = "BNSD";

  std::vector<gert::StorageShape> output_shapes(1, {{1, 100, 1024, 128}, {1, 100, 1024, 128}});

  std::vector<void *> output_shapes_ref(1);
  for (size_t i = 0; i < output_shapes.size(); ++i)
  {
    output_shapes_ref[i] = &output_shapes[i];
  }

  string compile_info_string = R"({
        "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
                          "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true, "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
                          "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288,
                          "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
                          "CORE_NUM": 40}
                          })";
  map<string, string> soc_infos;
  map<string, string> aicore_spec;
  map<string, string> intrinsics;
  map<string, string> socversions = {{"Short_SoC_version", "Ascend910B"}};
  GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);

  // platform info
  fe::PlatFormInfos platform_info;
  platform_info.Init();
  // compile info
  struct PromptFlashAttentionCompileInfo
  {
    uint32_t aivNum;
    uint32_t aicNum;
    uint64_t ubSize;
    uint64_t l1Size;
    uint64_t l0cSize;
    uint64_t l0ASize;
    uint64_t l0BSize;
    size_t defaultSysWorkspaceSize;
    platform_ascendc::SocVersion socShortName;
  };
  PromptFlashAttentionCompileInfo compile_info;

  std::string op_type("PromptFlashAttention");
  ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
  auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
  auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;

  // tilingParseFunc simulate
  auto kernel_holder = gert::KernelRunContextFaker()
                           .KernelIONum(2, 1)
                           .Inputs({const_cast<char *>(compile_info_string.c_str()), reinterpret_cast<void *>(&platform_info)})
                           .Outputs({&compile_info})
                           .Build();

  ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);

  ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);

  // tilingFunc simulate
  auto param = gert::TilingData::CreateCap(4096);
  auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
  auto ws_size = reinterpret_cast<gert::ContinuousVector *>(workspace_size_holer.get());
  ASSERT_NE(param, nullptr);
  auto holder = gert::TilingContextFaker()
                    .NodeIoNum(12, 1)
                    .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
                    .InputShapes({&query_shape, &key_shape, &value_shape, nullptr, nullptr, nullptr, nullptr, &deq_scale1, &quant_scale1, &deq_scale2, &quant_scale2, &quant_offset2})
                    .OutputShapes(output_shapes_ref)
                    .CompileInfo(&compile_info)
                    .PlatformInfo(reinterpret_cast<char *>(&platform_info))
                    .NodeInputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(2, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(3, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(4, ge::DT_BOOL, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(5, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(6, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(7, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(8, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(9, ge::DT_UINT64, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(10, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeInputTd(11, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeOutputTd(0, ge::DT_INT8, ge::FORMAT_ND, ge::FORMAT_ND)
                    .NodeAttrs({{"num_heads", ge::AnyValue::CreateFrom<int64_t>(num_heads)},
                                {"scale_value", ge::AnyValue::CreateFrom<float>(1.0)},
                                {"pre_tokens", ge::AnyValue::CreateFrom<int64_t>(pre_tokens)},
                                {"next_tokens", ge::AnyValue::CreateFrom<int64_t>(next_tokens)},
                                {"input_layout", ge::AnyValue::CreateFrom<string>(input_layout)},
                                {"num_key_value_heads", ge::AnyValue::CreateFrom<int64_t>(num_key_value_heads)}})
                    .TilingData(param.get())
                    .Workspace(ws_size)
                    .Build();

  gert::TilingContext *tiling_context = holder.GetContext<gert::TilingContext>();
  ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
  holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
  tiling_context->GetPlatformInfo()->SetPlatformRes("version", socversions); // label:"version" res:socversions

  // workspaces nullptr return failed
  EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_FAILED);
  // check tiling result
  auto tiling_key = tiling_context->GetTilingKey();
//   ASSERT_EQ(tiling_key, 1000000000000021012);
}